#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
直接测试数据提取功能
"""

import pandas as pd
import json
import os

def direct_test():
    """直接测试数据提取功能"""
    
    print("=== 直接测试数据提取功能 ===")
    
    # 检查文件存在性
    data_file = '整理后的行业指标数据.xlsx'
    dict_file = '行业指标字典表.xlsx'
    
    if not os.path.exists(data_file):
        print(f"❌ 数据文件不存在: {data_file}")
        return
    
    if not os.path.exists(dict_file):
        print(f"⚠️  字典文件不存在: {dict_file}")
    
    try:
        # 读取数据
        print("📊 正在读取数据文件...")
        df_data = pd.read_excel(data_file, sheet_name='整理后数据')
        print(f"✅ 数据文件读取成功: {len(df_data)} 条记录")
        print(f"列名: {list(df_data.columns)}")
        
        # 读取字典
        if os.path.exists(dict_file):
            print("📖 正在读取字典文件...")
            df_dict = pd.read_excel(dict_file)
            print(f"✅ 字典文件读取成功: {len(df_dict)} 个指标")
            print(f"字典列名: {list(df_dict.columns)}")
        else:
            df_dict = None
        
        # 显示基本统计信息
        print(f"\n=== 数据统计 ===")
        print(f"行业数量: {df_data['行业名称'].nunique()}")
        print(f"指标数量: {df_data['指标名称'].nunique()}")
        print(f"年份范围: {df_data['年份'].min()} - {df_data['年份'].max()}")
        
        # 显示前几个行业
        print(f"\n=== 行业列表（前10个）===")
        industries = sorted(df_data['行业名称'].unique())
        for i, industry in enumerate(industries[:10]):
            print(f"{i+1:2d}. {industry}")
        if len(industries) > 10:
            print(f"    ... 以及其他 {len(industries)-10} 个行业")
        
        # 显示年份
        print(f"\n=== 可用年份 ===")
        years = sorted(df_data['年份'].unique())
        print(f"年份: {years}")
        
        # 测试筛选功能
        print(f"\n=== 测试筛选功能 ===")
        
        # 测试1: 搜索包含"酒"的行业
        test_keyword = "酒"
        mask = df_data['行业名称'].str.contains(test_keyword, case=False, na=False)
        matching_industries = df_data[mask]['行业名称'].unique()
        print(f"包含'{test_keyword}'的行业: {list(matching_industries)}")
        
        # 测试2: 获取2023年数据
        if 2023 in df_data['年份'].values:
            data_2023 = df_data[df_data['年份'] == 2023]
            print(f"2023年数据: {len(data_2023)} 条记录")
        
        # 测试3: 组合筛选
        if len(matching_industries) > 0 and 2023 in df_data['年份'].values:
            combined = df_data[
                (df_data['行业名称'].str.contains(test_keyword, case=False, na=False)) &
                (df_data['年份'] == 2023)
            ]
            print(f"包含'{test_keyword}'的行业2023年数据: {len(combined)} 条记录")
            
            if len(combined) > 0:
                # 关联字典
                if df_dict is not None:
                    indicator_mapping = dict(zip(df_dict['指标名称'], df_dict['指标大类']))
                    combined_with_category = combined.copy()
                    combined_with_category['指标大类'] = combined_with_category['指标名称'].map(indicator_mapping)
                    combined_with_category['指标大类'] = combined_with_category['指标大类'].fillna('未分类')
                    
                    print(f"\n=== 示例输出数据 ===")
                    columns_order = ['行业名称', '年份', '指标大类', '指标名称', '优秀值', '良好值', '平均值', '较低值', '较差值']
                    sample_data = combined_with_category[columns_order].head(3)
                    
                    for i, (_, row) in enumerate(sample_data.iterrows()):
                        print(f"\n记录 {i+1}:")
                        for col in columns_order:
                            print(f"  {col}: {row[col]}")
                    
                    # 转换为JSON示例
                    sample_records = sample_data.to_dict('records')
                    for record in sample_records:
                        for key, value in record.items():
                            if pd.isna(value):
                                record[key] = None
                    
                    result = {
                        "query_info": {
                            "keyword": test_keyword,
                            "year": 2023,
                            "total_records": len(combined_with_category),
                            "industries": sorted(combined_with_category['行业名称'].unique().tolist()),
                            "indicators": sorted(combined_with_category['指标名称'].unique().tolist())
                        },
                        "data": sample_records
                    }
                    
                    print(f"\n=== JSON格式示例 ===")
                    json_output = json.dumps(result, ensure_ascii=False, indent=2)
                    print(json_output[:1500])  # 显示前1500个字符
                    if len(json_output) > 1500:
                        print("... (输出已截断)")
        
        print(f"\n✅ 测试完成")
        
    except Exception as e:
        print(f"❌ 测试过程中出错: {e}")
        import traceback
        traceback.print_exc()

if __name__ == "__main__":
    direct_test()
